# -*- coding: utf-8 -*-
"""
===============================
@Author     : Zuo WenTao
@Time       : 2024/11/8 11:15
@Description: 
@Change     : 
@File       : 判断是否为OpenAi向量.py
===============================
"""
import os


def list_directories(path):
    try:
        entries = os.listdir(path)

        directories = [entry for entry in entries if os.path.isdir(os.path.join(path, entry))]

        return directories
    except FileNotFoundError:
        print(f"目录 {path} 不存在")
        return []
    except PermissionError:
        print(f"没有权限访问目录 {path}")
        return []
    except Exception as e:
        print(f"发生错误: {e}")
        return []

def find_files_in_directory(directory, file_extension=None, max_count=None) -> list:
    """
    递归搜索指定目录下的所有文件
    :param directory: 指定的需要搜索的目录
    :param file_extension: 文件扩展名，如 ".txt"，".pdf" 等
    :param max_count: 需要获取的最大文件数量，如果为 None，则获取所有文件
    :return: 文件路径列表
    """
    files = []

    for root, dirs, filenames in os.walk(directory):
        for filename in filenames:
            file_path = os.path.join(root, filename)

            # 如果提供了文件类型筛选条件，则检查文件扩展名
            if file_extension is not None:
                if not filename.lower().endswith(file_extension.lower()):
                    continue

            files.append(file_path)

            # 如果达到最大文件数量限制，且 max_count 不为 None，则停止查找
            if max_count is not None and len(files) >= max_count:
                return files

    return files
# 示例使用
directory_path = "/mnt/journal04/Vector_data"
folders = list_directories(directory_path)

for folder in folders:
    if folder in ["cnki_book", "ChineseJournal_QC2", "ChineseJournal_QC", "chatjournal_向量", "chatfellows", "backup_test", "chinese_journal_all", "cnki"]: continue
    find_files_in_directory(directory=os.path.join(directory_path, folder), file_extension=".txt", max_count=10)
